import os
from typing import List
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

# --- 环境准备 ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码开始 ---

    # 1. 定义你的数据结构 (使用 Pydantic)
    # 我们创建一个类，它继承自 pydantic.BaseModel。
    # 这就是我们希望 AI 返回的数据的"蓝图"。
    class CharacterInfo(BaseModel):
        # Field 用于提供更详细的描述，这会帮助 AI 更好地理解我们想要什么。
        # 这里的 description 会被包含在给 AI 的指令中。
        name: str = Field(description="角色的名字")
        age: int = Field(description="角色的年龄")
        skills: List[str] = Field(description="角色的技能列表")

    # 2. 创建 Pydantic 解析器
    # 我们将刚刚定义的类传递给解析器。
    parser = PydanticOutputParser(pydantic_object=CharacterInfo)

    # 3. 创建提示模板
    # 这次，我们从解析器中获取格式化指令。
    # `parser.get_format_instructions()` 会自动生成一段告诉 AI 如何构建 JSON 的文本。
    prompt = ChatPromptTemplate.from_template(
        """
        你是一位顶级的用户信息分析师。
        请根据用户的描述，提取出关键的角色信息。
        {format_instructions}
        
        用户描述:
        {user_description}
        """
    )

    # 4. 创建链
    # 这一次，我们不再手动在提示里写 JSON 格式了。
    # 我们把从解析器生成的格式指令 `format_instructions` 注入到提示中。
    chain = prompt | ChatOpenAI(model=model_name) | parser

    # 5. 调用链
    user_input = "我介绍一下我的朋友，他叫李明，今年30岁了，是个非常厉害的程序员，尤其擅长 Python 和 Go 语言开发。"
    response = chain.invoke({
        "format_instructions": parser.get_format_instructions(),
        "user_description": user_input
    })

    print(f"用户输入: {user_input}")
    print(f"\nAI 提取的信息: {response}")
    print(f"结果的数据类型: {type(response)}") # 你会看到，结果是我们定义的 CharacterInfo 类型！
    
    # 6. 像操作对象一样访问数据！
    # 这就是 Pydantic 解析器的最大优势！
    print(f"\n角色姓名: {response.name}")
    print(f"角色年龄: {response.age}")
    print(f"角色技能: {response.skills}")
    # 还可以直接访问技能列表的第一个元素
    print(f"第一个技能: {response.skills[0]}")
    print("--------------------") 